# video_path = r"D:\home\code\videos_project\tennis.video"
video_path = r"D:\tennis\code\lable_video\temp\videos\1733989206751.h264"
import cv2
import numpy as np


def split_frame_data(frame_bts):
    pos = frame_bts.find("__DATETIME_SEP__".encode("utf-8"))
    if pos == -1:
        return None, None
    if pos !=0 :
        raise Exception("Invalid frame data")

    start_pos = pos + len("__DATETIME_SEP__".encode("utf-8"))
    end_pos = start_pos + 23 + 16
    timestamp = frame_bts[start_pos:end_pos].decode("utf-8")[0:23]
    jpeg_bytes = frame_bts[end_pos:]
    return timestamp, jpeg_bytes

# data = open("stroke.data", "r").readlines()

fp = open(video_path, "rb")
fp.seek(0)

chunk_size = 1024 * 100
last_frame_bts = b""
recent_frames = []


def get_frame_and_timestamp(fp):
    global last_frame_bts
    while True:
        pos = last_frame_bts.find("__FRAME_SEP__".encode("utf-8"))
        if pos == -1:
            bts = fp.read(chunk_size)
            if len(bts) == 0:
                raise Exception("EOF")
            last_frame_bts += bts
        else:
            cur_frame = last_frame_bts[0:pos]
            frame_timestamp, jpeg_bytes = split_frame_data(cur_frame)
            next_frame = last_frame_bts[pos+len("__FRAME_SEP__".encode("utf-8")):]
            last_frame_bts = next_frame
            return frame_timestamp, jpeg_bytes


def display_frame(frame_timestamp, jpeg_bytes):
        img = cv2.imdecode(np.frombuffer(jpeg_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        cv2.putText(img, str(frame_timestamp), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)  # Convert to string
        cv2.imshow("check", img)  # Also need to convert window name to string
        key = cv2.waitKey(0)  # Wait for user input (0 means wait indefinitely)
        if key & 0xFF == ord('q'):
            raise Exception("Quit")
        else:  # Enter key pressed
            pass

def timestamp_to_int(timestamp):
    # "2024-11-22 10:00:00.001"
    from datetime import datetime
    dt = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S.%f")
    # Convert to milliseconds since epoch (1970)
    return int(dt.timestamp() * 1000)


def save_frame(frame_timestamp, jpeg_bytes):
    with open(f"outputs/{frame_timestamp}.jpeg", "wb") as f:
        f.write(jpeg_bytes)


def save_label_result(label_result):
    with open("label.result", "w") as f:
        for timestamp, label in label_result.items():
            f.write(f"{timestamp} {label}\n")

label_result = {}
label_mapping = {
    "p": "Stroke",
    "y": "Serve",
    "l": "Land",
    "w": "Net"
}

def label_frame(recent_frames):
    current_index = len(recent_frames) - 1
    while current_index < len(recent_frames):
        timestamp, jpeg_bytes = recent_frames[current_index]
        img = cv2.imdecode(np.frombuffer(jpeg_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        cv2.putText(img, str(timestamp), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)  # Convert to string
        label_text = label_mapping[label_result[timestamp]] if timestamp in label_result else "None"
        cv2.putText(img, label_text, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)  # Convert to string
        cv2.imshow("check", img)  # Also need to convert window name to string
        key = cv2.waitKey(0)  # Wait for user input (0 means wait indefinitely)
        if key & 0xFF == ord('q'):
            # Quit
            raise Exception("Quit")
        elif key & 0xFF == ord('f'):
            # Forward to next frame
            current_index += 1
        elif key & 0xFF == ord('n'):
            # Remove label
            timestamp, _ = recent_frames[current_index]
            if timestamp in label_result:
                label_result.pop(timestamp)
        elif key & 0xFF == ord('p'):
            # Stroke
            timestamp, _ = recent_frames[current_index]
            label_result[timestamp] = "p"
        elif key & 0xFF == ord('y'):
            # Serve
            timestamp, _ = recent_frames[current_index]
            label_result[timestamp] = "y"
        elif key & 0xFF == ord('l'):
            # Land
            timestamp, _ = recent_frames[current_index]
            label_result[timestamp] = "l"
        elif key & 0xFF == ord('b'):
            # Back to previous frame
            current_index -= 1
            if current_index < 0:
                current_index = 0
        elif key & 0xFF == ord('w'):
            # Net
            timestamp, _ = recent_frames[current_index]
            label_result[timestamp] = "w"

        save_label_result(label_result)




        

frame_timestamp = -1


init_timestamp = -1
for line in open("label.result", "r").readlines():
    timestamp, label = line.strip().split(" ")
    timestamp = int(timestamp)
    label_result[timestamp] = label
    if init_timestamp < timestamp:
        init_timestamp = timestamp

while True:
    frame_timestamp, jpeg_bytes = get_frame_and_timestamp(fp)
    if frame_timestamp is None:
        continue
    frame_timestamp = timestamp_to_int(frame_timestamp)
    recent_frames.append((frame_timestamp, jpeg_bytes))
    if len(recent_frames) > 500:
        recent_frames.pop(0)
    save_frame(frame_timestamp, jpeg_bytes)
    # label_frame(recent_frames)












